import numpy as np
import common.Codec as codec
import struct
import tensorflow as tf
import convert.LabelProcess as lp
import convert.ImagePretreatment as ip


class ConvertFile:

    def __init__(self, root, full_file, head_length, new_size):
        self.__root = root
        self.__full_file = full_file
        self.__head_length = head_length
        self.__new_size = new_size

    def __read_file_into_memory(self, file, num, part_char_set):
        writer = tf.python_io.TFRecordWriter(
            "{0:s}/{1:d}.tfrecords".format(self.__root, num))
        for image, tag in self.__decode_hex(file):
            # decode label
            tag = struct.pack('<H', tag).decode('gb2312')
            # judge model
            if part_char_set != 'all':
                if (part_char_set.find(tag) == -1):
                    continue
                else:
                    tag = part_char_set.index(tag)
            else:
                # label to int
                tag = lp.char_to_int(tag)
            # resize image
            image = ip.image_resize(image, self.__new_size)
            # image = image.flatten().tolist()
            image = image.tostring()
            features = {"label": tf.train.Feature(
                int64_list=tf.train.Int64List(value=[tag])),
                "image": tf.train.Feature(
                    bytes_list=tf.train.BytesList(value=[image]))}
            # features = {"label": tf.train.Feature(
            #     int64_list=tf.train.Int64List(value=[tag]))}
            example = tf.train.Example(
                features=tf.train.Features(feature=features))
            writer.write(record=example.SerializeToString())
        writer.close()

    def __decode_hex(self, file):
        all_content = np.fromfile(file, dtype='uint8', count=-1)
        if all_content is None:
            print("error. head is null")
            return
        current = 0
        content_len = len(all_content)

        while current < content_len:
            # small
            sample_size = codec.hex2uint32(all_content[current:current + 4])
            # print('sample_size is {0:d}'.format(sample_size))
            tag_code = codec.hex2uint16(all_content[current + 4:current + 6])
            pic_width = codec.hex2uint16(all_content[current + 6:current + 8])
            pic_height = codec.hex2uint16(all_content[current + 8:current + 10])
            current = current + self.__head_length
            image_size = pic_height * pic_width
            # validate
            if sample_size != self.__head_length + image_size:
                print("error. decode error")
                return
            image = np.array(all_content[current:current + image_size]) \
                .reshape(pic_height, pic_width)
            current = current + image_size
            # show image
            # plt.figure()
            # plt.imshow(image, cmap='gray')
            # plt.show()
            yield image, tag_code

    def run(self, part_char_set):
        i = 1
        for file in self.__full_file:
            # read into memory
            # decode label
            # label change to number
            # resize image
            self.__read_file_into_memory(file, i, part_char_set)
            i = i + 1
